import tensorflow as tf
from matplotlib import pyplot as plt
from keras.datasets.cifar10 import load_data
from keras.preprocessing.image import ImageDataGenerator

(x_train, y_train), (x_test, y_test) = load_data()
x_train = x_train / 255
x_test = x_test / 255

image_data_generator = ImageDataGenerator(rotation_range=40,
                                          width_shift_range=0.2,
                                          height_shift_range=0.2,
                                          shear_range=20,
                                          zoom_range=0.2,
                                          horizontal_flip=True,
                                          vertical_flip=True)

index = 0
original_image = x_train[index]
augmented_image = image_data_generator.flow(x_train, shuffle=False).next()

plt.subplot(1, 2, 1)
plt.imshow(original_image)
plt.title('Original Image')
plt.subplot(1, 2, 2)
plt.imshow(augmented_image[0])
plt.title('Augmented Image')
plt.show()
